package tree;

import tree.operators.*;
import utils.Utils;

public class Tree 
{
	private IOp operator;
	private Tree parent;
	private Tree[] sons = new Tree[0];
	private int depth;
	private int numNodes, numFunctions, numTerminals;
	private double fitness;
	
	
	
	//private constructor used by the clone() method
	protected Tree() {}
	
	public Tree(int maxDepth)
	{
		this(null, maxDepth, 0);
	}
	
	protected Tree(Tree parent, int maxDepth, int currDepth)
	{
		this.parent = parent;
		this.depth = currDepth;
		
		if (currDepth == maxDepth)
		{
			operator = Op.createRandomTerminalOp();
		}
		else
		{
			operator = Op.createRandomOp();
		}
		
		initCounters();
		
		sons = new Tree[operator.arity()];
		for (int i = 0; i < sons.length; i++)
		{
			sons[i] = new Tree(this, maxDepth, currDepth + 1);
			initNumberOfNodes(i);
		}
	}

	private void initCounters() {
		numNodes = 1;
		if(operator.arity() == 0)
		{
			numFunctions = 0;
			numTerminals = 1;
		}
		else
		{
			numFunctions = 1;
			numTerminals = 0;
		}
	}
	
	public int getNumNodes()
	{
		return numNodes;
	}
	
	/**
	 * Evaluate the output of the function on the parameter x
	 * @param x - the input to the function
	 * @return
	 */
	public double evaluate(double x)
	{
		//prepare array of input arguments
		double[] args = new double[operator.arity()];
		for (int i = 0; i < args.length; i++)
		{
			args[i] = sons[i].evaluate(x);
		}
		return operator.evaluate(x, args);
		//return operator.evaluate(x);
	}
	
	/**
	 * Create a clone of this tree without a parent and with copies of the children
	 */
	public Tree clone()
	{
		return clone(0);
	}
	
	private Tree clone(int currDepth)
	{
		Tree clone = new Tree();
		clone.operator = this.operator; //it can point to the same operator, no need to create a new one
		
		//clone.numNodes = 1;
		clone.initCounters();
		clone.depth = currDepth;
		if (sons.length > 0)
		{
			clone.sons = new Tree[this.sons.length];

			for (int i = 0; i < sons.length; i++)
			{
				clone.sons[i] = this.sons[i].clone(currDepth+1);
				clone.sons[i].parent = clone;
				//clone.numNodes += clone.sons[i].numNodes;
				clone.initNumberOfNodes(i);
			}			
		}
		
		return clone;	
	}
	
	/**
	 * Get a random node in this tree (for selecting a splitting point)
	 * @return
	 * @throws Exception 
	 */
	public Tree getRandomNodeUniformDist() throws Exception
	{
		//return getNode(Utils.rand.nextInt(numNodes), 0, Integer.MAX_VALUE);
		int index = Utils.rand.nextInt(numNodes);
		Tree node = getNode(new IntWrapper(index), 0, Integer.MAX_VALUE);
		if(node == null)
		{
			return getNode(new IntWrapper(index), 0, Integer.MAX_VALUE);
		}
		return node;
	}
	
	/**
	 * Get a random node in this tree (for selecting a splitting point)
	 * @return
	 * @throws Exception 
	 */
	public Tree getRandomNodeKozaDist() throws Exception
	{
		double functionOrTerminal = Utils.rand.nextDouble();
		Tree node = null;
		if(functionOrTerminal < 0.9 && numFunctions > 0)
		{
			node = getNode(new IntWrapper(Utils.rand.nextInt(numFunctions)), 1, Integer.MAX_VALUE);
		}
		else
		{
			node = getNode(new IntWrapper(Utils.rand.nextInt(numTerminals)), 0, 0);
		}
		
		return node;
	}

	private Tree getNode(IntWrapper nodeIndex, int minArity, int maxArity)
	{
		if(this.operator.arity() >= minArity && this.operator.arity() <= maxArity)
		{
			if (nodeIndex.intValue == 0)
			{
				return this;
			}
			else
			{
				nodeIndex.intValue--;
			}
		}
		
		Tree foundNode = null;
		for(Tree son : sons)
		{
			foundNode = son.getNode(nodeIndex, minArity, maxArity);
			if(foundNode != null)
			{
				return foundNode;
			}
			
		}
		
		return null;
	}
	
	private Tree getFunction(int nodeIndex) throws Exception
	{
		if (nodeIndex == 0)
		{
			return this;
		}
		else if (nodeIndex - 1 < sons[0].numFunctions && sons[0].operator.arity() > 0)
		{
			return sons[0].getFunction(nodeIndex - 1);
		}
		else if(sons[0].operator.arity() > 0)
		{
			return sons[1].getFunction(nodeIndex - 1 - sons[0].numFunctions);
		}
		
		throw new Exception("BUG");
	}
	
	private Tree getTerminal(int nodeIndex) throws Exception
	{
		if(sons[0].operator.arity() == 0)
		{
			//Terminal
			if (nodeIndex == 0)
			{
				return this;
			}
			else if (nodeIndex - 1 < sons[0].numTerminals)
			{
				return sons[0].getTerminal(nodeIndex - 1);
			}
			else
			{
				return sons[1].getTerminal(nodeIndex - 1 - sons[0].numTerminals);
			}
		}
		
		if (nodeIndex == 0 && sons[0].operator.arity() == 0)
		{
			return this;
		}
		else if (nodeIndex - 1 < sons[0].numFunctions && sons[0].operator.arity() > 0)
		{
			return sons[0].getFunction(nodeIndex - 1);
		}
		else if(sons[0].operator.arity() > 0)
		{
			return sons[1].getFunction(nodeIndex - 1 - sons[0].numFunctions);
		}
		
		throw new Exception("BUG");
	}

	public String toString()
	{
		switch(this.sons.length)
		{
		case 0:
			return this.operator.toString();
		case 1:
			return this.operator.toString() + "(" + sons[0].toString() + ")";
		case 2:
			return "(" + sons[0].toString() + " " + this.operator.toString() + " " + sons[1].toString() + ")";
		default:
			return null;
		}
	}

	public void setFitness(double fitness) {
		this.fitness = fitness;
	}

	public double getFitness() {
		return fitness;
	}
	
	/**
	 * Attach a branch to this tree in a random location
	 * @param branch
	 * @throws Exception 
	 */
	private void attachBranch(Tree branch) throws Exception
	{
		//int nodeIdxToRemove = Utils.rand.nextInt(this.numNodes);
		
		//Tree nodeToRemove = getRandomNodeUniformDist();
		Tree nodeToRemove = getRandomNodeKozaDist();
		
		nodeToRemove.replace(branch);
		
		reevaluateSizeAndDepth(0);
		
	}
	
	private void reevaluateSizeAndDepth(int currDepth) 
	{
		this.depth = currDepth;
		
		initCounters();
		for (int i = 0; i < this.sons.length; i++)
		{
			sons[i].reevaluateSizeAndDepth(currDepth+1);
			initNumberOfNodes(i);
		}		
	}

	private void initNumberOfNodes(int childNum) {
		this.numNodes += sons[childNum].getNumNodes();
		this.numFunctions += sons[childNum].numFunctions;
		this.numTerminals += sons[childNum].numTerminals;
	}
	
	/**
	 * Replace this tree with the given tree
	 * size and depth will not be updated!!! - need to call reevaluateSizeAndDepth on the 
	 * root of the new tree
	 * @param branch
	 */
	private void replace(Tree branch)
	{
		if (parent == null)
		{
			this.operator = branch.operator;
			this.sons = branch.sons;
			for(Tree son : this.sons)
			{
				son.parent = this;
			}
		}
		else
		{
			branch.parent = this.parent;
			
			int numSonsOfParent = this.parent.sons.length;
			
			for (int i = 0; i < numSonsOfParent; i++)
			{
				if (parent.sons[i] == this)
				{
					parent.sons[i] = branch;
					break;
				}
			}
		}
	}
	
	public void mutate() throws Exception
	{		
		Tree cutoffNode = getRandomNodeUniformDist();
		int maxDepth = (int)Math.ceil(Math.log(cutoffNode.numNodes) / Math.log(2));
		Tree branch = new Tree(null, maxDepth, 0);
		cutoffNode.replace(branch);	
		reevaluateSizeAndDepth(0);
	}
	
	/**
	 * recombine two trees by selecting a random cutoff point in each one and
	 * attaching a copy of the first branch instead of the second branch
	 * @param t1
	 * @param t2
	 * @return
	 * @throws Exception 
	 */
	public static Tree recombine(Tree t1, Tree t2) throws Exception
	{
		//select random branch from tree 1
		//Tree branch = t1.getRandomNodeUniformDist().clone();
		Tree branch = t1.getRandomNodeKozaDist().clone();
		
		//create copy of tree 2
		Tree newTree = t2.clone();
		
		//attach branch to the clone
		newTree.attachBranch(branch);
		
		return newTree;
	}
	
	private class IntWrapper
	{

		public int intValue;
		
		public IntWrapper(int value) {
			this.intValue = value;
		}
	}
	
}
